{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\left\\{\\begin{array}{l}R(s)=\\left\\{\\begin{array}{l}+1,\\ s=(4,3)\\\\-1,\\ s=(4,2)\\\\-0.02\\ ,\\ otherwise.\\end{array}\\right.\\\\\\\n",
       "    \\gamma=0.99\\\\\\\n",
       "    正常走概率0.8, 向左向右0.1\\\n",
       "    \\end{array}\\right.$"
      ],
      "text/plain": [
       "<IPython.core.display.Latex object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Latex\n",
    "from IPython.display import display, Math, Latex\n",
    "print_latex = lambda latex_str: display(Latex(latex_str))\n",
    "\n",
    "print_latex(r'$\\left\\{\\begin{array}{l}R(s)=\\left\\{\\begin{array}{l}+1,\\ s=(4,3)\\\\-1,\\ s=(4,2)\\\\-0.02\\ ,\\ otherwise.\\end{array}\\right.\\\\\\\n",
    "    \\gamma=0.99\\\\\\\n",
    "    正常走概率0.8, 向左向右0.1\\\n",
    "    \\end{array}\\right.$')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "38\n",
      "[[ 0.8664417   0.90710179  0.94380762  1.        ]\n",
      " [ 0.83070053        -inf  0.7916229  -0.14387273]\n",
      " [ 0.79110902  0.75630717  0.74271026  0.61485363]]\n",
      "[['r' 'r' 'r' 'Stop']\n",
      " ['u' None 'u' 'u']\n",
      " ['u' 'l' 'u' 'l']]\n"
     ]
    }
   ],
   "source": [
    "# 价值迭代算法\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import pprint\n",
    "\n",
    "max_x = 3\n",
    "max_y = 2\n",
    "min_x = 0\n",
    "min_y = 0\n",
    "direc = ['u', 'r', 'd', 'l']\n",
    "a_size = len(direc)\n",
    "offset = {'u': (0, 1), 'd': (0, -1), 'l': (-1, 0), 'r': (1, 0)}\n",
    "p_normal = 0.8\n",
    "p_abnormal = 0.1\n",
    "R_rest = -0.02\n",
    "gamma = 0.99\n",
    "e = 1e-30\n",
    "obstacle = {(1, 1)}\n",
    "states = {(x, y) for x in range(min_x, max_x+1) for y in range(min_y, max_y+1)} - obstacle\n",
    "next_state = {}\n",
    "for s in states:\n",
    "    if s[0] != 3 or s[1] != 2:\n",
    "        for a in direc:\n",
    "            if a == 'u' and s[1] == max_y or a == 'd' and s[1] == min_y or a == 'l' and s[0] == min_x or a == 'r' and s[0] == max_x:\n",
    "                next_state[s + (a,)] = s\n",
    "            else: \n",
    "                if tuple(np.add(s, offset[a])) in obstacle:\n",
    "                    next_state[s + (a,)] = s\n",
    "                else:\n",
    "                    next_state[s + (a,)] = tuple(np.add(s, offset[a]))\n",
    "    else:\n",
    "        for a in direc:\n",
    "            next_state[s + (a,)] = s\n",
    "\n",
    "            \n",
    "# Reward Function\n",
    "R = np.ones((max_x-min_x+1, max_y-min_y+1)) * R_rest\n",
    "R[3][2] = 1\n",
    "R[3][1] = -1\n",
    "Rs = {s: R[s[0]][s[1]] for s in states}\n",
    "Rs.update({s: 0 for s in obstacle})\n",
    "\n",
    "# Value Function\n",
    "V = np.zeros((max_x-min_x+1, max_y-min_y+1))\n",
    "Vs = {s: V[s[0]][s[1]] for s in states}\n",
    "Vs.update({s: -np.inf for s in obstacle})\n",
    "\n",
    "# Get Value Function\n",
    "n = 0\n",
    "while True:\n",
    "    lastV = pd.Series(Vs)\n",
    "    for s in states:\n",
    "        max_next = -np.inf\n",
    "        if s == (3, 2):\n",
    "            max_next = 0\n",
    "        else:\n",
    "            for i in range(len(direc)):\n",
    "                cur_next = p_normal * Vs[next_state[s + (direc[i],)]] \\\n",
    "                            + p_abnormal*Vs[next_state[s + (direc[(i+1)%a_size],)]] \\\n",
    "                            + p_abnormal*Vs[next_state[s + (direc[(i+a_size-1)%a_size],)]]\n",
    "                if cur_next > max_next:\n",
    "                    max_next = cur_next\n",
    "        Vs[s] = Rs[s] + gamma*max_next\n",
    "    V = pd.Series(Vs)\n",
    "    n += 1\n",
    "    if np.sum(np.square(lastV-V)) < e:\n",
    "        break\n",
    "    \n",
    "# Get Policy\n",
    "policy = {}\n",
    "for s in states:\n",
    "        max_next = -np.inf\n",
    "        best_policy = direc[0]\n",
    "        if s == (3, 2):\n",
    "            policy[s] = 'Stop'\n",
    "        else:\n",
    "            for i in range(len(direc)):\n",
    "                cur_next = p_normal * Vs[next_state[s + (direc[i],)]] \\\n",
    "                            + p_abnormal*Vs[next_state[s + (direc[(i+1)%a_size],)]] \\\n",
    "                            + p_abnormal*Vs[next_state[s + (direc[(i+a_size-1)%a_size],)]]\n",
    "                if cur_next > max_next:\n",
    "                    best_policy = direc[i]\n",
    "                    max_next = cur_next\n",
    "            policy[s] = best_policy\n",
    "\n",
    "print(n)\n",
    "V = np.array([[Vs[(x, y)] for x in range(min_x, max_x+1)] for y in range(min_y, max_y+1)])[::-1]\n",
    "print(V)\n",
    "Po = np.array([[policy[(x, y)] if (x, y) not in obstacle else None for x in range(min_x, max_x+1)] for y in range(min_y, max_y+1)])[::-1]\n",
    "print(Po)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 价值迭代算法\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import pprint\n",
    "\n",
    "max_x = 3\n",
    "max_y = 2\n",
    "min_x = 0\n",
    "min_y = 0\n",
    "direc = ['u', 'r', 'd', 'l']\n",
    "a_size = len(direc)\n",
    "offset = {'u': (0, 1), 'd': (0, -1), 'l': (-1, 0), 'r': (1, 0)}\n",
    "p_normal = 0.8\n",
    "p_abnormal = 0.1\n",
    "R_rest = -0.02\n",
    "gamma = 0.99\n",
    "e = 1e-30\n",
    "obstacle = {(1, 1)}\n",
    "states = {(x, y) for x in range(min_x, max_x+1) for y in range(min_y, max_y+1)} - obstacle\n",
    "next_state = {}\n",
    "for s in states:\n",
    "    if s[0] != 3 or s[1] != 2:\n",
    "        for a in direc:\n",
    "            if a == 'u' and s[1] == max_y or a == 'd' and s[1] == min_y or a == 'l' and s[0] == min_x or a == 'r' and s[0] == max_x:\n",
    "                next_state[s + (a,)] = s\n",
    "            else: \n",
    "                if tuple(np.add(s, offset[a])) in obstacle:\n",
    "                    next_state[s + (a,)] = s\n",
    "                else:\n",
    "                    next_state[s + (a,)] = tuple(np.add(s, offset[a]))\n",
    "    else:\n",
    "        for a in direc:\n",
    "            next_state[s + (a,)] = s\n",
    "\n",
    "            \n",
    "# Reward Function\n",
    "R = np.ones((max_x-min_x+1, max_y-min_y+1)) * R_rest\n",
    "R[3][2] = 1\n",
    "R[3][1] = -1\n",
    "Rs = {s: R[s[0]][s[1]] for s in states}\n",
    "Rs.update({s: 0 for s in obstacle})\n",
    "\n",
    "# Value Function\n",
    "V = np.zeros((max_x-min_x+1, max_y-min_y+1))\n",
    "Vs = {s: V[s[0]][s[1]] for s in states}\n",
    "Vs.update({s: -np.inf for s in obstacle})\n",
    "\n",
    "# Get Value Function\n",
    "n = 0\n",
    "while True:\n",
    "    lastV = pd.Series(Vs)\n",
    "    for s in states:\n",
    "        max_next = -np.inf\n",
    "        if s == (3, 2):\n",
    "            max_next = 0\n",
    "        else:\n",
    "            for i in range(len(direc)):\n",
    "                cur_next = p_normal * Vs[next_state[s + (direc[i],)]] \\\n",
    "                            + p_abnormal*Vs[next_state[s + (direc[(i+1)%a_size],)]] \\\n",
    "                            + p_abnormal*Vs[next_state[s + (direc[(i+a_size-1)%a_size],)]]\n",
    "                if cur_next > max_next:\n",
    "                    max_next = cur_next\n",
    "        Vs[s] = Rs[s] + gamma*max_next\n",
    "    V = pd.Series(Vs)\n",
    "    n += 1\n",
    "    if np.sum(np.square(lastV-V)) < e:\n",
    "        break\n",
    "    \n",
    "# Get Policy\n",
    "policy = {}\n",
    "for s in states:\n",
    "        max_next = -np.inf\n",
    "        best_policy = direc[0]\n",
    "        if s == (3, 2):\n",
    "            policy[s] = 'Stop'\n",
    "        else:\n",
    "            for i in range(len(direc)):\n",
    "                cur_next = p_normal * Vs[next_state[s + (direc[i],)]] \\\n",
    "                            + p_abnormal*Vs[next_state[s + (direc[(i+1)%a_size],)]] \\\n",
    "                            + p_abnormal*Vs[next_state[s + (direc[(i+a_size-1)%a_size],)]]\n",
    "                if cur_next > max_next:\n",
    "                    best_policy = direc[i]\n",
    "                    max_next = cur_next\n",
    "            policy[s] = best_policy\n",
    "\n",
    "print(n)\n",
    "V = np.array([[Vs[(x, y)] for x in range(min_x, max_x+1)] for y in range(min_y, max_y+1)])[::-1]\n",
    "print(V)\n",
    "Po = np.array([[policy[(x, y)] if (x, y) not in obstacle else None for x in range(min_x, max_x+1)] for y in range(min_y, max_y+1)])[::-1]\n",
    "print(Po)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
